use hierr::Error;
use hiopt::{options, raw_args_from_i8, OptErr, Options};
use hisym::*;

static OPTS: Options<'static> = options![
    "o:", "output:", "l:", "lib:", "s", "show", "h", "help", "c", "compress", "i", "e", "a", "d", "m", "r"
];

struct Context<'a> {
    target: &'a str,
    out: &'a str,
    lib: bool,
    show: bool,
    show_import: bool,
    show_export: bool,
    show_dyn: bool,
    show_sym: bool,
    show_rel: bool,
    help: bool,
    compress: bool,
}

impl Context<'_> {
    fn null() -> Self {
        Self {
            target: "",
            out: "",
            lib: false,
            show: false,
            show_import: false,
            show_export: false,
            show_dyn: false,
            show_sym: false,
            show_rel: false,
            help: false,
            compress: false,
        }
    }

    fn new() -> Self {
        let mut ctx = Self::null();
        let args = unsafe { raw_args_from_i8(hictor::args()) };
        for opt in OPTS.opt_iter(args) {
            let Ok((idx, arg)) = opt else {
                print_err(opt.err().unwrap());
                help();
            };
            match idx {
                0..=1 => ctx.out = arg.unwrap(),
                2..=3 => ctx.lib = true,
                4..=5 => ctx.show = true,
                6..=7 => ctx.help = true,
                8..=9 => ctx.compress = true,
                10 => ctx.show_import = true,
                11 => ctx.show_export = true,
                13 => ctx.show_dyn = true,
                14 => ctx.show_sym = true,
                15 => ctx.show_rel = true,
                12 => {
                    ctx.show_import = true;
                    ctx.show_export = true;
                    ctx.show_dyn = true;
                    ctx.show_sym = true;
                    ctx.show_rel = true;
                }
                _ => unreachable!(),
            }
        }

        let mut iter = OPTS.noopt_iter(&args[1..]);
        if let Some(target) = iter.next() {
            ctx.target = target;
        }
        if ctx.target.is_empty() || iter.next().is_some() {
            help();
        }
        ctx
    }
}

fn print_err(err: OptErr) {
    match err {
        OptErr::NoShortOpt(f) => eprintln!("unsupport option: {}", f as char),
        OptErr::NoLongOpt(f) => eprintln!("unsupport option: {f}"),
        OptErr::NoArg(n) => eprintln!("option {} requires an argument", OPTS[n].name()),
    }
}

fn help() -> ! {
    let args = unsafe { raw_args_from_i8(hictor::args()) };
    let program = OPTS.noopt_iter(args).next().unwrap();
    println!("{program} [options] <target file> ");
    println!("options:");
    println!("-o <output file name> | --output <output file name>");
    println!("-l <dynamic library name> | --lib <dynamic library name>: compress the symbols in the target file that depend on the dynamic library, this option can appear more than once");
    println!("-s | --show: display information about each segment in the target file");
    println!("-i: display import symbol if -s is set");
    println!("-e: display export symbol if -s is set");
    println!("-a: display import/export symbol if -s is set");
    println!("-h | --help: display help information");
    unsafe { libc::exit(-1) };
}

fn main() {
    let ctx = Context::new();
    if ctx.help {
        help();
    }
    if ctx.target.is_empty() {
        return;
    }

    let fd = match ElfFd::open(ctx.target) {
        Ok(fd) => fd,
        Err(err) => {
            eprintln!("{err:?}");
            return;
        }
    };

    if fd.is_bits::<u64>() {
        if ctx.show {
            show::<u64>(&ctx, fd);
            return;
        }
        if !ctx.compress {
            return;
        }
        if !ctx.lib {
            compress_export::<u64>(&ctx, fd);
        } else {
            compress_import::<u64>(&ctx, fd);
        }
    } else {
        if ctx.show {
            show::<u32>(&ctx, fd);
            return;
        }
        if !ctx.compress {
            return;
        }
        if !ctx.lib {
            compress_export::<u32>(&ctx, fd);
        } else {
            compress_import::<u32>(&ctx, fd);
        }
    }
}

fn show<T: ElfBits>(ctx: &Context<'_>, fd: ElfFd) {
    let file = ElfDynFile::<T>::open(fd).unwrap();
    println!(
        "{:<2} {:<20} {:<20} {:<10} {:<10} {:<10} {:<5} {:<10}",
        "No", "name", "type", "offset", "size", "addr", "align", "flags",
    );
    for (n, hdr) in file.section_iter::<T>().enumerate() {
        let name = file.read_section_name::<T>(hdr.sh_name()).unwrap_or("");
        let tname = sh_type_name(hdr.sh_type());
        let offset = hdr.sh_offset();
        let size = hdr.sh_size();
        let addr = hdr.sh_addr();
        let align = hdr.sh_addralign();
        let flags = sh_flags_name(hdr.sh_flags());
        println!("{n:<2} {name:<20} {tname:<20} {offset:<10x} {size:<10x} {addr:<10x} {align:<5x} {flags:<10}");
    }

    println!(
        "{:<2} {:<20} {:<10} {:<10} {:<10} {:<10} {:<10} {:<5}",
        "No", "type", "offset", "vaddr", "paddr", "filesz", "memsz", "align"
    );
    for (n, hdr) in file.segment_iter::<T>().enumerate() {
        let tname = ph_type_name(hdr.ph_type());
        let offset = hdr.ph_offset();
        let vaddr = hdr.ph_vaddr();
        let paddr = hdr.ph_paddr();
        let filesz = hdr.ph_filesz();
        let memsz = hdr.ph_memsz();
        let align = hdr.ph_align();
        let addr = if align > 0 { (offset + align - 1) & !(align - 1) } else { offset };
        println!(
            "{n:<2} {tname:<20} {offset:<10x} {vaddr:<10x} {paddr:<10x} {filesz:<10x} {memsz:<10x} {align:<5x} {addr:<10x}"
        );
    }

    if ctx.show_import {
        for (symndx, symname) in file.import_iter() {
            println!("import {symndx}.{symname}");
        }
    }

    if ctx.show_export {
        for (symndx, symname) in file.export_iter() {
            println!("export {symndx}.{symname}");
        }
    }

    if ctx.show_dyn {
        for e in file.dyn_iter::<T>() {
            println!("dyn {} {:x}", d_tag_is_ptr(e.d_tag()), e.d_val_ptr());
        }
    }

    if ctx.show_sym {
        for e in file.symbol_iter::<T>() {
            println!("addr {:x}", e.st_value());
        }
    }

    if ctx.show_rel {
        for e in file.rel_iter::<T>() {
            println!("r_off {:x}", e.r_offset());
        }
        for e in file.rela_iter::<T>() {
            println!("r_off {:x}", e.r_offset());
        }
    }
}

fn compress_export<T: ElfBits>(ctx: &Context<'_>, fd: ElfFd) {
    println!("compress export symbol...");
    let mut new = [0_u8; 32];
    let mut file = ElfDynFile::<T>::open(fd).unwrap();
    let nbuckets = file.hash_buckets();
    let mut buff = file.new_name_buff().unwrap();
    for (symndx, name) in file.export_iter() {
        let new_name = compress_name(name, &mut new);
        if new_name.len() >= name.len() {
            continue;
        }
        buff.write(nbuckets, symndx, new_name);
    }
    let diff = file.update_merge(&mut buff);
    for (symndx, name) in file.export_iter() {
        assert!(Some(symndx) == file.find(name));
    }
    if write_output(ctx, file.as_bytes()) {
        println!("{diff} bytes are saved");
    }
}

fn compress_import<T: ElfBits>(ctx: &Context<'_>, fd: ElfFd) {
    println!("compress import symbol");
    let mut file = ElfDynFile::<T>::open(fd).unwrap();
    let mut buff = file.new_name_buff().unwrap();

    let args = unsafe { raw_args_from_i8(hictor::args()) };
    for opt in OPTS.opt_iter(args) {
        let Ok((idx, arg)) = opt else {
            continue;
        };
        if idx != 2 && idx != 3 {
            continue;
        }
        let lib = arg.unwrap();
        let Ok(libfd) = ElfFd::open(lib) else {
            eprintln!("failed to open library");
            return;
        };
        let Ok(libfile) = ElfDynFile::<T>::open(libfd) else {
            eprintln!("failed to open library");
            return;
        };
        compress_import_with::<T>(&file, &mut buff, libfile);
    }
    let diff = file.update_merge(&mut buff);
    if write_output(ctx, file.as_bytes()) {
        println!("{diff} bytes are saved");
    }
}

fn compress_import_with<T: ElfBits>(file: &ElfDynFile<T>, buff: &mut NameBuff, lib: ElfDynFile<T>) {
    let mut new = [0_u8; 32];
    let nbuckets = lib.hash_buckets();
    for (symndx, symname) in file.import_iter() {
        let new_name = buff.name_by_symndx(symndx).unwrap();
        if new_name.len() < symname.len() {
            continue;
        }

        let new_name = compress_name(symname, &mut new);
        // lib中的符号有可能已经被压缩过一次了
        if lib.find(symname).is_none() {
            if !buff.write(nbuckets, symndx, new_name) {
                continue;
            }
            //最终压缩的名字可能在用户提供的名字后增加后缀
            let new_name = buff.name_by_symndx(symndx).unwrap();
            if lib.find(new_name).is_none() {
                //压缩前后的名字都没有找到，则放弃
                buff.restore(symndx, symname);
            }
            continue;
        }
        buff.write(nbuckets, symndx, new_name);
    }
}

fn write_output(ctx: &Context<'_>, bytes: &[u8]) -> bool {
    //ctx.output 来自命令行输入，是一个C字符串格式
    let fd = unsafe {
        libc::open(
            ctx.out.as_ptr().cast::<libc::c_char>(),
            libc::O_CREAT | libc::O_TRUNC | libc::O_RDWR,
            0o777,
        )
    };
    if fd == -1 {
        eprintln!("failed to create file - {}", ctx.out);
        return false;
    }
    let size = unsafe { libc::write(fd, bytes.as_ptr().cast::<libc::c_void>(), bytes.len()) };
    if size == -1 {
        eprintln!("failed to write data: {:?}", Error::last_error());
    } else if (size as usize) < bytes.len() {
        eprintln!(
            "NoSpace to write all data, write {} bytes, but return {} bytes",
            bytes.len(),
            size
        );
    }

    unsafe { libc::close(fd) };
    (size as usize) == bytes.len()
}

fn compress_name<'a, 'b, 'c>(name: &'a str, new: &'b mut [u8; 32]) -> &'c str
where
    'a: 'c,
    'b: 'c,
{
    if name.len() <= new.len() {
        return name;
    }

    let prefix_len = 10;
    new[..prefix_len].copy_from_slice(&name.as_bytes()[..prefix_len]);
    new[0] = b'_';
    new[1] = b'C';

    for (b, out) in md5::compute(name)[0..8].iter().zip(new[10..].chunks_mut(2)) {
        unsafe { print_u8(*b, out) };
    }
    unsafe { core::str::from_utf8_unchecked(&new[0..26]) }
}

// out.len() >= 2
unsafe fn print_u8(b: u8, out: &mut [u8]) {
    let bytes = b"0123456789ABCDEF";
    out[0] = bytes[(b >> 4) as usize];
    out[1] = bytes[(b & 0x0F) as usize];
}
